import numpy as np
import torch
import matplotlib.pyplot as plt

fig = plt.figure()
# 构建x,y点图
x1 = torch.linspace(0, 1, 20)
w1 = torch.Tensor([2.])
w2 = torch.Tensor([-1.])
b1 = torch.Tensor([4.])
b2 = torch.Tensor([1.])
x2 = w1 * x1 ** 2 + b1
y = torch.cos(x2 * w2) + b2
noise = torch.normal(0, 0.3, size=y.shape)
y_noise = y + noise

ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)
ax1.plot(x1.detach().numpy(), y.detach().numpy(), 'ro-')
ax1.plot(x1.detach().numpy(), y_noise.detach().numpy(), 'gv')

# 训练
w1_predict = torch.Tensor([0.1])
w2_predict = torch.Tensor([0.2])
b1_predict = torch.Tensor([0.3])
b2_predict = torch.Tensor([0.4])
w1_predict.requires_grad = True
w2_predict.requires_grad = True
b1_predict.requires_grad = True
b2_predict.requires_grad = True
result = torch.cos((w1_predict * x1 ** 2 + b1_predict) * w2_predict) + b2_predict
line, = ax1.plot(x1.detach().numpy(), result.detach().numpy(), 'b--')

epochs = 1000
lr = 0.1
line_x = [1]
line_y = [1]
line2, = ax2.plot(line_x, line_y, 'b-')
for epoch in range(epochs):
    y_predict = 1 / torch.exp(w2_predict * (1 / torch.exp(w1_predict * x1 ** 2) + b1_predict)) + b2_predict
    e = torch.mean((y_noise - y_predict) ** 2)
    e.backward()

    with torch.no_grad():
        for param in [w2_predict, b2_predict, w1_predict, b1_predict]:
            slope = param.grad
            param -= lr * slope
            param.grad.zero_()

    if epoch % 10 == 0:
        print(f"epoch {epoch}/{epochs} ---- loss:{e}")
        line_y.append(e.item())
        line_x.append(len(line_y))

        line2.set_data(line_x, line_y)
        ax2.relim()
        ax2.autoscale_view()

    line.set_data(x1.detach().numpy(), y_predict.detach().numpy())
    plt.pause(0.1)

# plt.show()
